import os
import logging
import pandas as pd
import torch
import time
import data_augment
from tqdm import tqdm  # 添加tqdm库用于进度显示
from transformers import BertTokenizer
from bert import Bert
from train_eval import train
from utils import config_to_json_string, read_csv, get_dateset_labels


class ErnieConfig:
    """ERNIE模型训练配置类，存储所有训练相关参数"""

    def __init__(self):
        # 预训练模型路径配置
        self.pretrain_path = '../pretrain_models/ERNIE'
        _config_file = 'bert_config.json'  # 模型配置文件
        _model_file = 'pytorch_model.bin'  # 预训练权重文件
        _tokenizer_file = 'vocab.txt'  # 分词器词汇表
        self.config_file = os.path.join(self.pretrain_path, _config_file)
        self.model_name_or_path = os.path.join(self.pretrain_path, _model_file)
        self.tokenizer_file = os.path.join(self.pretrain_path, _tokenizer_file)

        # 数据路径配置
        self.train_path = '../data/KUAKE/KUAKE-QQR_train.json'  # 训练数据路径
        self.dev_path = '../data/KUAKE/KUAKE-QQR_dev.json'  # 验证数据路径
        self.test_path = '../data/KUAKE/KUAKE-QQR_test.json'  # 测试数据路径
        self.aug_data_path = '../data/KUAKE/KUAKE-QQR_augment.csv'  # 增强数据路径

        # 模型和任务配置
        self.use_model = 'bert'  # 使用的基础模型架构
        self.task = 'KUAKE'  # 任务名称
        self.models_name = 'ernie'  # 模型名称
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 计算设备
        self.requires_grad = True  # 是否需要计算梯度
        self.class_list = ['0', '1', '2']  # 分类标签列表
        self.num_labels = len(self.class_list)  # 分类类别数量

        # 数据统计信息
        self.train_num_examples = 0  # 训练样本数量
        self.dev_num_examples = 0  # 验证样本数量
        self.test_num_examples = 0  # 测试样本数量

        # 模型超参数
        self.hidden_dropout_prob = 0.1  # 隐藏层dropout概率
        self.multi_drop = 5  # 多次dropout的次数
        self.hidden_size = 768  # 隐藏层维度
        self.early_stop = False  # 是否启用早停
        self.require_improvement = 1000  # 早停判断的步数

        # 训练超参数
        self.num_train_epochs = 10  # 训练轮数
        self.batch_size = 128  # 批次大小
        self.pad_size = 64  # 输入序列最大长度
        self.learning_rate = 1e-5  # 基础学习率
        self.head_learning_rate = 1e-3  # 分类头学习率(通常高于基础模型)
        self.weight_decay = 0.01  # 权重衰减系数
        self.warmup_proportion = 0.1  # warmup步数比例

        # 日志配置
        self.is_logging2file = True  # 是否将日志写入文件
        self.logging_dir = '../logging' + '/' + self.models_name  # 日志目录

        # 模型保存配置
        self.save_path = '../my_model'  # 模型保存路径
        self.save_file = self.models_name  # 模型保存文件名

        # 数据增强配置
        self.data_augment = False  # 是否使用数据增强
        self.diff_learning_rate = False  # 是否使用差分学习率
        self.stop_word_valid = True  # 是否使用停用词过滤

        # 进度显示配置
        self.show_progress = True  # 是否显示进度条


if __name__ == '__main__':
    # 初始化配置
    config = ErnieConfig()
    print(f"训练设备: {config.device}")

    # 数据标签统计(需查看数据分布可取消注释)
    labels = get_dateset_labels(config.train_path)
    label0 = labels.count(0)  # 8946
    label1 = labels.count(1)  # 2603
    label2 = labels.count(2)  # 3451

    # 配置日志
    logging_filename = None
    if config.is_logging2file is True:
        file = time.strftime('%Y-%m-%d_%H-%M-%S') + '.log'  # 日志文件名格式
        logging_filename = os.path.join(config.logging_dir, file)
        if not os.path.exists(config.logging_dir):
            os.makedirs(config.logging_dir)  # 创建日志目录

    # 设置日志格式和级别
    logging.basicConfig(filename=logging_filename, format='%(levelname)s: %(message)s', level=logging.INFO)

    # 初始化分词器
    tokenizer = BertTokenizer.from_pretrained(config.tokenizer_file)

    # 加载训练数据
    print("正在加载训练数据...")
    train_examples = read_csv(config.train_path)
    config.train_num_examples = len(train_examples)
    print(f"训练数据加载完成，共{config.train_num_examples}条样本")

    # 数据增强处理
    if config.data_augment:
        # 加载增强数据并合并到训练集
        print("正在加载增强数据...")
        augment_examples = data_augment.dataframe_to_list(pd.read_csv(config.aug_data_path))
        train_examples.extend(augment_examples)
        logging.info(
            f"使用数据增强，原始数据 {len(train_examples) - len(augment_examples)} 条，增强后 {len(train_examples)} 条")
        print(f"数据增强完成，训练集共有{len(train_examples)}条样本")

    # 加载验证数据
    print("正在加载验证数据...")
    dev_examples = read_csv(config.dev_path)
    config.dev_num_examples = len(dev_examples)
    print(f"验证数据加载完成，共{config.dev_num_examples}条样本")

    # 初始化模型
    model = Bert(config)
    # print(model)  # 如需查看模型结构可取消注释

    # 记录配置信息到日志
    logging.info(f"模型配置: {config_to_json_string(config)}")

    # 计算总训练步数
    num_train_steps = int(len(train_examples) / config.batch_size * config.num_train_epochs)
    print(f"总训练步数: {num_train_steps}")

    # 开始训练
    train(
        config=config,
        model=model,
        tokenizer=tokenizer,
        train_data=train_examples,
        dev_data=dev_examples,
    )